package predictionEngine;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

public class RandomForest {
    private int size;
    private List<TreeNode> trees;
    private double outOfBagError=0;
    private HashMap<double[],Double> predictions,predictionCount;
    public RandomForest(int size) {
        this.size=size;
        trees = new ArrayList<TreeNode>(size);
    }
    
    
    public void createForest(Dataset dataset){
        createForest(dataset, false);
    }
        
    public void createForest(Dataset dataset,Boolean calcOutOfBag){
        int consideredAttributes = (int)Math.round(Math.sqrt(dataset.getNumberOfAttributes()));
        System.out.println("Considering " + consideredAttributes +" randomly selected attributes");
        for(int i = 0; i < size; i++){
            Bootstrap strap = new Bootstrap(dataset);
            TreeNode tree = new TreeNode(null,strap.getBag(),consideredAttributes);
            trees.add(tree);
            if(calcOutOfBag){
                addOutOfBagPredictions(tree,strap);
            }
        }
        if(calcOutOfBag){
            outOfBagError=calcOutOfBagMAE(dataset);
            System.out.println("Out of bag MAE: " + outOfBagError);
        }
    }
    
    public double getOutOfBagMAE(){
        return outOfBagError;
    }
    
    private double calcOutOfBagMAE(Dataset dataset){
        double MAE = 0;
        for(double[] instance:dataset.getInstances()){
            double prediction = predictions.get(instance)/predictionCount.get(instance);
            MAE+=Math.abs(prediction-instance[instance.length-1]);
        }
        return MAE/dataset.size();
    }
    
    private void addOutOfBagPredictions(TreeNode tree,Bootstrap strap){
        if(predictions==null){
            predictions = new HashMap<double[],Double>(strap.getBag().size());
            predictionCount = new HashMap<double[],Double>(strap.getBag().size());
        }
        for(double[] instance:strap.getOutofBag().getInstances()){
            Double count = predictionCount.get(instance);
            Double prediction = predictions.get(instance);
            if(count==null){
                count = new Double(0);
                prediction = new Double(0);
            }
            count++;            
            prediction+=tree.execute(instance);  
            
            predictionCount.put(instance, count);
            predictions.put(instance, prediction);
        }        
    }
    
    public double execute(double[] instance){
        double prediction=0;
        for(TreeNode tree:trees)
            prediction+=tree.execute(instance);
        return prediction/size;
    }
}
